use std::collections::HashMap;

#[derive(Debug, Clone, PartialEq)]
pub struct Bytecode {
    pub instructions: Vec<Instruction>,
    pub constants: Vec<i64>,
    pub functions: HashMap<String, usize>, // 函数名到指令索引的映射
}

impl Bytecode {
    pub fn new() -> Self {
        Bytecode {
            instructions: Vec::new(),
            constants: Vec::new(),
            functions: HashMap::new(),
        }
    }

    pub fn add_constant(&mut self, value: i64) -> usize {
        self.constants.push(value);
        self.constants.len() - 1
    }

    pub fn add_function(&mut self, name: String, start_index: usize) {
        self.functions.insert(name, start_index);
    }

    pub fn get_function(&self, name: &str) -> Option<usize> {
        self.functions.get(name).copied()
    }
}

#[derive(Debug, Clone, PartialEq)]
pub enum Instruction {
    PushInt(i64),
    PushConst(usize), // 使用常量池索引
    Pop,
    Add,
    Sub,
    Mul,
    Div,
    Load(usize), // 加载局部变量
    Store(usize), // 存储局部变量
    Call(String), // 调用函数
    NativeCall(String), // 原生函数调用
    Ret,
    Jump(usize), // 无条件跳转
    JumpIfZero(usize), // 条件跳转
    Halt,
    Print,
    Dup,
    Swap,
    Eq,
    Ne,
    Lt,
    Gt,
    Le,
    Ge,
    And,
    Or,
    Not,
    Neg,
    Inc,
    Dec,
    PushTrue,
    PushFalse,
}

pub struct BytecodeGenerator {
    bytecode: Bytecode,
}

impl BytecodeGenerator {
    pub fn new() -> Self {
        BytecodeGenerator {
            bytecode: Bytecode::new(),
        }
    }

    pub fn emit(&mut self, instruction: Instruction) {
        self.bytecode.instructions.push(instruction);
    }

    pub fn add_constant(&mut self, value: i64) -> usize {
        self.bytecode.add_constant(value)
    }

    pub fn add_function(&mut self, name: String, start_index: usize) {
        self.bytecode.add_function(name, start_index);
    }

    pub fn finalize(self) -> Bytecode {
        self.bytecode
    }

    pub fn get_current_index(&self) -> usize {
        self.bytecode.instructions.len()
    }
}

impl Bytecode {
    pub fn to_bytes(&self) -> Vec<u8> {
        let mut bytes = Vec::new();
        
        // 写入常量数量
        bytes.extend_from_slice(&(self.constants.len() as u32).to_le_bytes());
        
        // 写入常量
        for &constant in &self.constants {
            bytes.extend_from_slice(&constant.to_le_bytes());
        }
        
        // 写入指令数量
        bytes.extend_from_slice(&(self.instructions.len() as u32).to_le_bytes());
        
        // 写入指令
        for instruction in &self.instructions {
            match instruction {
                Instruction::PushInt(val) => {
                    bytes.push(0x01);
                    bytes.extend_from_slice(&val.to_le_bytes());
                }
                Instruction::PushConst(idx) => {
                    bytes.push(0x02);
                    bytes.extend_from_slice(&(*idx as u32).to_le_bytes());
                }
                Instruction::Pop => bytes.push(0x03),
                Instruction::Add => bytes.push(0x04),
                Instruction::Sub => bytes.push(0x05),
                Instruction::Mul => bytes.push(0x06),
                Instruction::Div => bytes.push(0x07),
                Instruction::Load(idx) => {
                    bytes.push(0x08);
                    bytes.extend_from_slice(&(*idx as u32).to_le_bytes());
                }
                Instruction::Store(idx) => {
                    bytes.push(0x09);
                    bytes.extend_from_slice(&(*idx as u32).to_le_bytes());
                }
                Instruction::Call(name) => {
                    bytes.push(0x0A);
                    let name_bytes = name.as_bytes();
                    bytes.extend_from_slice(&(name_bytes.len() as u32).to_le_bytes());
                    bytes.extend_from_slice(name_bytes);
                }
                Instruction::NativeCall(name) => {
                    bytes.push(0x20);
                    let name_bytes = name.as_bytes();
                    bytes.extend_from_slice(&(name_bytes.len() as u32).to_le_bytes());
                    bytes.extend_from_slice(name_bytes);
                }
                Instruction::Ret => bytes.push(0x0B),
                Instruction::Jump(offset) => {
                    bytes.push(0x0C);
                    bytes.extend_from_slice(&(*offset as u32).to_le_bytes());
                }
                Instruction::JumpIfZero(offset) => {
                    bytes.push(0x0D);
                    bytes.extend_from_slice(&(*offset as u32).to_le_bytes());
                }
                Instruction::Halt => bytes.push(0x0E),
                Instruction::Print => bytes.push(0x0F),
                Instruction::Dup => bytes.push(0x10),
                Instruction::Swap => bytes.push(0x11),
                Instruction::Eq => bytes.push(0x12),
                Instruction::Ne => bytes.push(0x13),
                Instruction::Lt => bytes.push(0x14),
                Instruction::Gt => bytes.push(0x15),
                Instruction::Le => bytes.push(0x16),
                Instruction::Ge => bytes.push(0x17),
                Instruction::And => bytes.push(0x18),
                Instruction::Or => bytes.push(0x19),
                Instruction::Not => bytes.push(0x1A),
                Instruction::Neg => bytes.push(0x1B),
                Instruction::Inc => bytes.push(0x1C),
                Instruction::Dec => bytes.push(0x1D),
                Instruction::PushTrue => bytes.push(0x1E),
                Instruction::PushFalse => bytes.push(0x1F),
            }
        }
        
        bytes
    }

    pub fn from_bytes(bytes: &[u8]) -> Result<Self, String> {
        let mut cursor = 0;
        
        if bytes.len() < 4 {
            return Err("Invalid bytecode: too short".to_string());
        }
        
        let mut bytecode = Bytecode::new();
        
        // 读取常量
        let const_count = u32::from_le_bytes([bytes[cursor], bytes[cursor+1], bytes[cursor+2], bytes[cursor+3]]) as usize;
        cursor += 4;
        
        for _ in 0..const_count {
            if cursor + 8 > bytes.len() {
                return Err("Invalid bytecode: incomplete constant".to_string());
            }
            let val = i64::from_le_bytes([bytes[cursor], bytes[cursor+1], bytes[cursor+2], bytes[cursor+3], bytes[cursor+4], bytes[cursor+5], bytes[cursor+6], bytes[cursor+7]]);
            bytecode.constants.push(val);
            cursor += 8;
        }
        
        // 读取指令
        if cursor + 4 > bytes.len() {
            return Err("Invalid bytecode: missing instruction count".to_string());
        }
        
        let instr_count = u32::from_le_bytes([bytes[cursor], bytes[cursor+1], bytes[cursor+2], bytes[cursor+3]]) as usize;
        cursor += 4;
        
        for _ in 0..instr_count {
            if cursor >= bytes.len() {
                return Err("Invalid bytecode: incomplete instruction".to_string());
            }
            
            let opcode = bytes[cursor];
            cursor += 1;
            
            let instruction = match opcode {
                0x01 => {
                    if cursor + 8 > bytes.len() {
                        return Err("Invalid bytecode: incomplete PushInt".to_string());
                    }
                    let val = i64::from_le_bytes([bytes[cursor], bytes[cursor+1], bytes[cursor+2], bytes[cursor+3], bytes[cursor+4], bytes[cursor+5], bytes[cursor+6], bytes[cursor+7]]);
                    cursor += 8;
                    Instruction::PushInt(val)
                }
                0x02 => {
                    if cursor + 4 > bytes.len() {
                        return Err("Invalid bytecode: incomplete PushConst".to_string());
                    }
                    let idx = u32::from_le_bytes([bytes[cursor], bytes[cursor+1], bytes[cursor+2], bytes[cursor+3]]) as usize;
                    cursor += 4;
                    Instruction::PushConst(idx)
                }
                0x03 => Instruction::Pop,
                0x04 => Instruction::Add,
                0x05 => Instruction::Sub,
                0x06 => Instruction::Mul,
                0x07 => Instruction::Div,
                0x08 => {
                    if cursor + 4 > bytes.len() {
                        return Err("Invalid bytecode: incomplete Load".to_string());
                    }
                    let idx = u32::from_le_bytes([bytes[cursor], bytes[cursor+1], bytes[cursor+2], bytes[cursor+3]]) as usize;
                    cursor += 4;
                    Instruction::Load(idx)
                }
                0x09 => {
                    if cursor + 4 > bytes.len() {
                        return Err("Invalid bytecode: incomplete Store".to_string());
                    }
                    let idx = u32::from_le_bytes([bytes[cursor], bytes[cursor+1], bytes[cursor+2], bytes[cursor+3]]) as usize;
                    cursor += 4;
                    Instruction::Store(idx)
                }
                0x0A => {
                    if cursor + 4 > bytes.len() {
                        return Err("Invalid bytecode: incomplete Call".to_string());
                    }
                    let name_len = u32::from_le_bytes([bytes[cursor], bytes[cursor+1], bytes[cursor+2], bytes[cursor+3]]) as usize;
                    cursor += 4;
                    
                    if cursor + name_len > bytes.len() {
                        return Err("Invalid bytecode: incomplete Call name".to_string());
                    }
                    
                    let name = String::from_utf8(bytes[cursor..cursor+name_len].to_vec())
                        .map_err(|_| "Invalid UTF-8 in function name")?;
                    cursor += name_len;
                    Instruction::Call(name)
                }
                0x0B => Instruction::Ret,
                0x0C => {
                    if cursor + 4 > bytes.len() {
                        return Err("Invalid bytecode: incomplete Jump".to_string());
                    }
                    let offset = u32::from_le_bytes([bytes[cursor], bytes[cursor+1], bytes[cursor+2], bytes[cursor+3]]) as usize;
                    cursor += 4;
                    Instruction::Jump(offset)
                }
                0x0D => {
                    if cursor + 4 > bytes.len() {
                        return Err("Invalid bytecode: incomplete JumpIfZero".to_string());
                    }
                    let offset = u32::from_le_bytes([bytes[cursor], bytes[cursor+1], bytes[cursor+2], bytes[cursor+3]]) as usize;
                    cursor += 4;
                    Instruction::JumpIfZero(offset)
                }
                0x0E => Instruction::Halt,
                0x0F => Instruction::Print,
                0x10 => Instruction::Dup,
                0x11 => Instruction::Swap,
                0x12 => Instruction::Eq,
                0x13 => Instruction::Ne,
                0x14 => Instruction::Lt,
                0x15 => Instruction::Gt,
                0x16 => Instruction::Le,
                0x17 => Instruction::Ge,
                0x18 => Instruction::And,
                0x19 => Instruction::Or,
                0x1A => Instruction::Not,
                0x1B => Instruction::Neg,
                0x1C => Instruction::Inc,
                0x1D => Instruction::Dec,
                0x1E => Instruction::PushTrue,
                0x1F => Instruction::PushFalse,
                0x20 => {
                    if cursor + 4 > bytes.len() {
                        return Err("Invalid bytecode: incomplete NativeCall".to_string());
                    }
                    let name_len = u32::from_le_bytes([bytes[cursor], bytes[cursor+1], bytes[cursor+2], bytes[cursor+3]]) as usize;
                    cursor += 4;
                    
                    if cursor + name_len > bytes.len() {
                        return Err("Invalid bytecode: incomplete NativeCall name".to_string());
                    }
                    
                    let name = String::from_utf8(bytes[cursor..cursor+name_len].to_vec())
                        .map_err(|_| "Invalid UTF-8 in native function name")?;
                    cursor += name_len;
                    Instruction::NativeCall(name)
                }
                _ => return Err(format!("Unknown opcode: {}", opcode)),
            };
            
            bytecode.instructions.push(instruction);
        }
        
        Ok(bytecode)
    }
}